import os
import time
import random
import json
from tqdm import tqdm  # 进度条工具（当前代码未直接使用，预留接口）

import torch
import torch.nn as nn
from torch.nn.utils import clip_grad_norm_  # 梯度裁剪工具，防止梯度爆炸
from torch.utils.data import DataLoader  # 数据加载器
from tensorboardX import SummaryWriter  # TensorBoard日志记录，用于可视化
from torch.optim.lr_scheduler import StepLR, MultiStepLR  # 学习率调度器

import numpy as np
from configs.opts import parser  # 命令行参数解析器
from model.main_model import supv_main_model as main_model  # 有监督主模型
from utils import AverageMeter, Prepare_logger, get_and_save_args  # 工具函数/类
from utils.Recorder import Recorder  # 实验记录器（代码备份、日志保存）
from dataset.AVE_dataset import AVEDataset  # AVE数据集类
import torch.nn.functional as F  # PyTorch函数式接口

# =================================  随机种子配置 ============================
# 固定随机种子以确保实验可复现
SEED = 43
random.seed(SEED)
np.random.seed(seed=SEED)
torch.manual_seed(seed=SEED)
torch.cuda.manual_seed(seed=SEED)
torch.backends.cudnn.deterministic = True  # 禁用cudnn的非确定性算法
torch.backends.cudnn.benchmark = False  # 禁用benchmark模式（影响速度但保证可复现）

# 加载模型配置文件（JSON格式）
config_path = 'configs/main.json'
with open(config_path) as fp:
    config = json.load(fp)
print(config)  # 打印配置信息，确认加载正确
# =============================================================================

def AVPSLoss(av_simm, soft_label):
    """
    音视频对相似度损失函数（完全有监督场景）
    对应论文中的公式(8, 9)，用于优化音视频模态间的相似度匹配
    Args:
        av_simm: 音视频相似度分数 (batch_size, seq_len)
        soft_label: 软标签（用于相似度的目标匹配值）
    Returns:
        均方误差损失值
    """
    relu_av_simm = F.relu(av_simm)  # ReLU激活，过滤负相似度
    sum_av_simm = torch.sum(relu_av_simm, dim=-1, keepdim=True)  # 时序维度求和
    avg_av_simm = relu_av_simm / (sum_av_simm + 1e-8)  # 归一化（加1e-8避免除零）
    loss = nn.MSELoss()(avg_av_simm, soft_label)  # 均方误差损失
    return loss


def main():
    # 全局变量声明（跨函数共享关键参数和统计信息）
    global args, logger, writer, dataset_configs
    global best_accuracy, best_accuracy_epoch
    best_accuracy, best_accuracy_epoch = 0, 0  # 记录最佳准确率及对应epoch

    # 加载并合并配置参数（命令行参数覆盖默认配置）
    dataset_configs = get_and_save_args(parser)
    parser.set_defaults(**dataset_configs)
    args = parser.parse_args()

    # GPU设备配置（指定使用的GPU编号）
    os.environ['CUDA_DEVICE_ORDER'] = "PCI_BUS_ID"
    os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu

    '''创建实验快照目录：用于备份代码和保存模型'''
    if not os.path.exists(args.snapshot_pref):
        os.makedirs(args.snapshot_pref)
    # 若从 checkpoint 恢复训练，使用 checkpoint 所在目录作为快照目录
    if os.path.isfile(args.resume):
        args.snapshot_pref = os.path.dirname(args.resume)

    # 初始化日志记录器（区分训练和评估模式）
    logger = Prepare_logger(args, eval=args.evaluate)

    # 打印实验配置（训练模式下）
    if not args.evaluate:
        logger.info(f'\nCreating folder: {args.snapshot_pref}')
        logger.info('\nRuntime args\n\n{}\n'.format(json.dumps(vars(args), indent=4)))
    else:
        logger.info(f'\nLog file will be save in {args.snapshot_pref}/Eval.log.')

    '''数据集加载：初始化训练和测试数据加载器'''
    train_dataloader = DataLoader(
        AVEDataset('./data/', split='train'),  # 训练集
        batch_size=args.batch_size,  # 训练批次大小
        shuffle=True,  # 训练集打乱顺序
        num_workers=8,  # 数据加载线程数
        pin_memory=True  # 锁定内存，加速GPU数据传输
    )

    test_dataloader = DataLoader(
        AVEDataset('./data/', split='test'),  # 测试集
        batch_size=args.test_batch_size,  # 测试批次大小
        shuffle=False,  # 测试集不打乱
        num_workers=8,
        pin_memory=True
    )

    '''模型配置：初始化模型、优化器、调度器和损失函数'''
    mainModel = main_model(config['model'])  # 初始化有监督主模型
    mainModel = nn.DataParallel(mainModel).cuda()  # 多GPU并行训练（移至GPU）
    learned_parameters = mainModel.parameters()  # 待优化的模型参数
    optimizer = torch.optim.Adam(learned_parameters, lr=args.lr)  # Adam优化器
    # 学习率调度器：多步衰减（在指定epoch降低学习率）
    scheduler = MultiStepLR(optimizer, milestones=[10, 20, 30], gamma=0.5)
    # 损失函数：二分类交叉熵（带logits）和多分类交叉熵
    criterion = nn.BCEWithLogitsLoss().cuda()  # 用于事件置信度预测
    criterion_event = nn.CrossEntropyLoss().cuda()  # 用于事件类别分类

    '''从checkpoint恢复训练/评估'''
    if os.path.isfile(args.resume):
        logger.info(f"\nLoading Checkpoint: {args.resume}\n")
        mainModel.load_state_dict(torch.load(args.resume))  # 加载模型权重
    elif args.resume != "" and (not os.path.isfile(args.resume)):
        raise FileNotFoundError  # 若指定resume但文件不存在，抛出异常

    '''仅执行评估（不训练）'''
    if args.evaluate:
        logger.info(f"\nStart Evaluation..")
        validate_epoch(mainModel, test_dataloader, criterion, criterion_event, epoch=0, eval_only=True)
        return

    '''TensorBoard可视化和代码备份初始化'''
    writer = SummaryWriter(args.snapshot_pref)  # TensorBoard日志写入器
    recorder = Recorder(args.snapshot_pref, ignore_folder="Exps/")  # 实验记录器
    recorder.writeopt(args)  # 保存实验参数到文件

    '''训练与测试主循环'''
    for epoch in range(args.n_epoch):
        # 训练一个epoch，返回平均损失
        loss = train_epoch(mainModel, train_dataloader, criterion, criterion_event, optimizer, epoch)

        # 按指定频率评估模型（或最后一个epoch强制评估）
        if ((epoch + 1) % args.eval_freq == 0) or (epoch == args.n_epoch - 1):
            acc = validate_epoch(mainModel, test_dataloader, criterion, criterion_event, epoch)
            # 若当前准确率优于历史最佳，保存模型 checkpoint
            if acc > best_accuracy:
                best_accuracy = acc
                best_accuracy_epoch = epoch
                save_checkpoint(
                    mainModel.state_dict(),
                    top1=best_accuracy,
                    task='Supervised',
                    epoch=epoch + 1,
                )
            print("-----------------------------")
            print("best acc and epoch:", best_accuracy, best_accuracy_epoch)
            print("-----------------------------")
        # 更新学习率
        scheduler.step()


def train_epoch(model, train_dataloader, criterion, criterion_event, optimizer, epoch):
    """
    单轮训练函数：遍历训练集，更新模型参数
    Args:
        model: 训练的模型
        train_dataloader: 训练数据加载器
        criterion: 二分类损失函数（事件置信度）
        criterion_event: 多分类损失函数（事件类别）
        optimizer: 优化器
        epoch: 当前epoch编号
    Returns:
        本轮训练的平均损失
    """
    # 初始化统计指标计数器
    batch_time = AverageMeter()  # 批次处理时间
    data_time = AverageMeter()   # 数据加载时间
    losses = AverageMeter()      # 损失值
    train_acc = AverageMeter()   # 训练准确率
    end_time = time.time()       # 计时起点

    model.train()  # 设置模型为训练模式（启用Dropout、BatchNorm训练模式）
    model.double()  # 转换模型为double精度（适配输入特征的数据类型）
    optimizer.zero_grad()  # 初始化梯度为0

    # 遍历训练集批次
    for n_iter, batch_data in enumerate(train_dataloader):
        # 更新数据加载时间
        data_time.update(time.time() - end_time)

        '''加载批次数据并预处理'''
        visual_feature, audio_feature, labels = batch_data  # 视觉特征、音频特征、标签
        labels = labels.double().cuda()  # 标签转换为double精度并移至GPU

        '''模型前向传播'''
        is_event_scores, event_scores, audio_visual_gate, av_score = model(visual_feature, audio_feature)
        # 调整输出维度以适配损失计算（时序维度转置）
        is_event_scores = is_event_scores.transpose(1, 0).squeeze().contiguous()
        audio_visual_gate = audio_visual_gate.transpose(1, 0).squeeze().contiguous()

        '''标签处理：提取前景标签、事件类别标签'''
        labels_foreground = labels[:, :, :-1]  # 去掉背景类标签（[batch, seq_len, 28]）
        labels_BCE, labels_evn = labels_foreground.max(-1)  # BCE标签（置信度）和时序事件类别
        labels_event, _ = labels_evn.max(-1)  # 视频级事件类别（时序维度取最大值）

        '''计算多部分损失并求和'''
        loss_is_event = criterion(is_event_scores, labels_BCE.cuda())  # 事件置信度损失
        label_is_gate = criterion(audio_visual_gate, labels_BCE.cuda())  # 模态门控置信度损失
        loss_cas = criterion_event(av_score, labels_event.cuda())  # 类别激活序列损失
        loss_event_class = criterion_event(event_scores, labels_event.cuda())  # 事件类别分类损失
        loss = loss_is_event + label_is_gate + loss_event_class + loss_cas  # 总损失

        '''反向传播计算梯度'''
        loss.backward()

        '''计算训练准确率并更新统计'''
        acc = compute_accuracy_supervised(is_event_scores, event_scores, labels)
        train_acc.update(acc.item(), visual_feature.size(0) * 10)  # 按样本数加权更新

        '''梯度裁剪：防止梯度爆炸'''
        if args.clip_gradient is not None:
            total_norm = clip_grad_norm_(model.parameters(), args.clip_gradient)

        '''优化器更新模型参数'''
        optimizer.step()
        optimizer.zero_grad()  # 重置梯度

        '''更新统计指标'''
        losses.update(loss.item(), visual_feature.size(0) * 10)  # 按时序步长加权损失
        batch_time.update(time.time() - end_time)
        end_time = time.time()

        '''记录TensorBoard日志（迭代级）'''
        writer.add_scalar('Train_data/loss', losses.val, epoch * len(train_dataloader) + n_iter + 1)

        '''打印终端日志（按指定频率）'''
        if n_iter % args.print_freq == 0:
            logger.info(
                f'Train Epoch: [{epoch}][{n_iter}/{len(train_dataloader)}]\t'
                f'Loss {losses.val:.4f} ({losses.avg:.4f})\t'
                f'Prec@1 {train_acc.val:.3f} ({train_acc.avg: .3f})'
            )

    '''记录TensorBoard日志（epoch级）'''
    writer.add_scalar('Train_epoch_data/epoch_loss', losses.avg, epoch)

    return losses.avg


@torch.no_grad()  # 禁用梯度计算，加速评估并节省内存
def validate_epoch(model, test_dataloader, criterion, criterion_event, epoch, eval_only=False):
    """
    单轮评估函数：遍历测试集，计算模型性能（准确率、损失）
    Args:
        model: 待评估模型
        test_dataloader: 测试数据加载器
        criterion: 二分类损失函数
        criterion_event: 多分类损失函数
        epoch: 当前epoch编号
        eval_only: 是否仅执行评估（不记录TensorBoard）
    Returns:
        测试集平均准确率
    """
    # 初始化统计指标计数器
    batch_time = AverageMeter()
    data_time = AverageMeter()
    losses = AverageMeter()
    accuracy = AverageMeter()
    end_time = time.time()

    model.eval()  # 设置模型为评估模式（禁用Dropout、固定BatchNorm）
    model.double()  # 转换为double精度

    # 遍历测试集批次
    for n_iter, batch_data in enumerate(test_dataloader):
        data_time.update(time.time() - end_time)

        '''加载批次数据并预处理'''
        visual_feature, audio_feature, labels = batch_data
        labels = labels.double().cuda()
        bs = visual_feature.size(0)  # 批次大小

        '''模型前向传播'''
        is_event_scores, event_scores, audio_visual_gate, _ = model(visual_feature, audio_feature)
        is_event_scores = is_event_scores.transpose(1, 0).squeeze()
        audio_visual_gate = audio_visual_gate.transpose(1, 0).squeeze()

        '''标签处理（与训练阶段一致）'''
        labels_foreground = labels[:, :, :-1]
        labels_BCE, labels_evn = labels_foreground.max(-1)
        labels_event, _ = labels_evn.max(-1)

        '''计算评估损失'''
        loss_is_event = criterion(is_event_scores, labels_BCE.cuda())
        loss_is_gate = criterion(audio_visual_gate, labels_BCE.cuda())
        loss_event_class = criterion_event(event_scores, labels_event.cuda())
        loss = loss_is_event + loss_event_class + loss_is_gate

        '''计算准确率并更新统计'''
        acc = compute_accuracy_supervised(is_event_scores, event_scores, labels)
        accuracy.update(acc.item(), bs * 10)

        '''更新统计指标'''
        batch_time.update(time.time() - end_time)
        end_time = time.time()
        losses.update(loss.item(), bs * 10)

        '''打印终端日志（按指定频率）'''
        if n_iter % args.print_freq == 0:
            logger.info(
                f'Test Epoch [{epoch}][{n_iter}/{len(test_dataloader)}]\t'
                f'Loss {losses.val:.4f} ({losses.avg:.4f})\t'
                f'Prec@1 {accuracy.val:.3f} ({accuracy.avg:.3f})'
            )

    '''记录TensorBoard日志（非仅评估模式）'''
    if not eval_only:
        writer.add_scalar('Val_epoch_data/epoch_loss', losses.avg, epoch)
        writer.add_scalar('Val_epoch/Accuracy', accuracy.avg, epoch)

    # 打印本轮评估的准确率
    logger.info(
        f'**************************************************************************\t'
        f"\tEvaluation results (acc): {accuracy.avg:.4f}%."
    )
    return accuracy.avg


def compute_accuracy_supervised(is_event_scores, event_scores, labels):
    """
    计算有监督场景下的准确率
    Args:
        is_event_scores: 事件置信度分数 (batch_size, seq_len)
        event_scores: 事件类别分数 (batch_size, num_classes)
        labels: 真实标签 (batch_size, seq_len, num_classes+1)
    Returns:
        准确率（百分比）
    """
    _, targets = labels.max(-1)  # 提取真实标签的类别索引（包含背景类）
    # 事件置信度分数sigmoid激活（转换为0-1概率）
    is_event_scores = is_event_scores.sigmoid()
    # 置信度阈值判断（>0.5视为事件）
    scores_pos_ind = is_event_scores > 0.5
    scores_mask = scores_pos_ind == 0  # 背景掩码（置信度<=0.5）
    # 事件类别预测（取分数最高的类别）
    _, event_class = event_scores.max(-1)
    # 构建预测结果：事件位置填充类别，背景位置填充背景索引（28）
    pred = scores_pos_ind.long()
    pred *= event_class[:, None]
    pred[scores_mask] = 28  # 28表示背景类
    # 计算正确预测数并求准确率
    correct = pred.eq(targets)
    correct_num = correct.sum().double()
    acc = correct_num * (100. / correct.numel())  # 转换为百分比
    return acc


def save_checkpoint(state_dict, top1, task, epoch):
    """
    保存模型checkpoint
    Args:
        state_dict: 模型状态字典
        top1: 当前最佳准确率
        task: 任务类型（用于文件名）
        epoch: 当前epoch编号
    """
    # 构建checkpoint文件名（包含epoch、准确率、任务类型）
    model_name = f'{args.snapshot_pref}/model_epoch_{epoch}_top1_{top1:.3f}_task_{task}_best_model.pth.tar'
    torch.save(state_dict, model_name)  # 保存模型权重


if __name__ == '__main__':
    main()  # 启动训练/评估流程